# utils/data_processing.py
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler


def load_data(data_path="data/glv_data.npz"):
    """加载s和p矩阵"""
    data = np.load(data_path)
    s = data["s"]  # 形状：(num_samples, N)
    p = data["p"]  # 形状：(num_samples, N)
    return s, p


def preprocess_data(s, p, test_size=0.2, random_state=42):
    """
    预处理数据：划分训练集和测试集，对输入标准化（可选）
    返回：训练集(s_train, p_train)、测试集(s_test, p_test)
    """
    # 划分训练集和测试集（由于样本少，测试集比例设小一点）
    s_train, s_test, p_train, p_test = train_test_split(
        s, p, test_size=test_size, random_state=random_state
    )

    # 注意：s是0-1向量，无需标准化；p是概率分布（和为1），也可不用标准化
    # 这里仅返回划分后的数据
    return s_train, s_test, p_train, p_test